import os
import numpy as np
from langchain_community.document_loaders import UnstructuredPDFLoader

print("=== 环境检查 ===")
print(f"NumPy 版本: {np.__version__}")
print(f"NumPy 测试: {np.array([1, 2, 3])}")

# # 设置缓存目录
# os.environ['UNSTRUCTURED_CACHE_DIR'] = os.path.expanduser('~/.cache/unstructured')
# os.makedirs(os.environ['UNSTRUCTURED_CACHE_DIR'], exist_ok=True)

print("\n=== 测试完整功能 ===")

# 测试1：带语言参数的hi_res模式
try:
    print("测试1: hi_res模式 + 中英文语言")
    pdf_documents = UnstructuredPDFLoader(
        "../knowledge_base/sample.pdf",
        mode="elements",
        strategy="hi_res",
        infer_table_structure=False,  # 先不启用表格识别
        languages=["eng", "chi_sim"],
    ).load()
    print(f"✅ 成功! 加载了 {len(pdf_documents)} 个元素")

    # 显示元素统计
    from collections import Counter

    categories = Counter([doc.metadata.get('category', 'Unknown') for doc in pdf_documents])
    print("元素类型统计:", dict(categories))

except Exception as e:
    print(f"❌ 失败: {e}")

# 测试2：启用表格识别
try:
    print("\n测试2: hi_res模式 + 表格识别")
    pdf_documents = UnstructuredPDFLoader(
        "../knowledge_base/sample.pdf",
        mode="elements",
        strategy="hi_res",
        infer_table_structure=True,  # 启用表格识别
        languages=["eng", "chi_sim"],
    ).load()
    print(f"✅ 表格识别成功! 加载了 {len(pdf_documents)} 个元素")

    # 检查表格元素
    table_elements = [doc for doc in pdf_documents if doc.metadata.get('category', '') == 'Table']
    print(f"找到 {len(table_elements)} 个表格元素")

    for i, table in enumerate(table_elements[:2]):  # 显示前两个表格
        print(f"\n表格 {i + 1}:")
        print(f"  文本内容: {table.page_content[:200]}...")
        if table.metadata.get('text_as_html'):
            print(f"  HTML表格预览: {table.metadata['text_as_html']}")
        else:
            print("  无HTML表格数据")

except Exception as e:
    print(f"❌ 表格识别失败: {e}")
    import traceback

    traceback.print_exc()

# 测试3：显示一些文本内容示例
try:
    print("\n测试3: 内容示例")
    pdf_documents = UnstructuredPDFLoader(
        "../knowledge_base/sample.pdf",
        mode="elements",
        strategy="hi_res",
        infer_table_structure=False,
        languages=["eng", "chi_sim"],
    ).load()

    # 显示不同类型的元素
    print("内容示例:")
    for i, doc in enumerate(pdf_documents[:10]):  # 前10个元素
        category = doc.metadata.get('category', '')
        content_preview = doc.page_content[:100].replace('\n', ' ') if doc.page_content else "空内容"
        print(f"  {i + 1}. [{category}] {content_preview}...")

except Exception as e:
    print(f"❌ 内容显示失败: {e}")
"""
修复原始 notebook 中的代码
"""
print("+"*20)
try:
    from langchain_community.document_loaders import UnstructuredPDFLoader

    pdf_documents = UnstructuredPDFLoader(
        "../knowledge_base/sample.pdf",
        mode="elements",
        strategy="hi_res",
        infer_table_structure=True,
        languages=["eng", "chi_sim"],
    ).load()
    print(pdf_documents)
    
except Exception as e:
    print("Error:", e)
    import traceback
    traceback.print_exc()
